# 服务器，用来控制训练过程
import copy
import torch 
import logging
# logging.basicConfig(level = logging.INFO,format = '%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logging.basicConfig(format ='%(message)s')
logger = logging.getLogger()
logger.setLevel(logging.INFO)


'''
约定上，服务器必须继承该类，避免使用强制接口。并实现以下方法；
1. __init__()初始化自己的客户端
2. aggragate_by_weight()本地特殊的聚合方式
3. start_federated() 启动联邦学习的函数

其他的函数。可以直接写在基类里作为通用的方法。
'''

class FedServerBase(object):
    def __init__(self,worker_list,model,server_params,client_params):
        """初始化服务器。并设置相关的参数。
        1. 初始化服务器相关的参数：device
        2. 初始化深度学习相关参数：model、optimizer
        3. 初始化联邦学习相关参数：rounds、ratio、inner_lr、outer_lr
        3. 完成客户端配置。主要调用客户端设置参数、处理数据。

        Args:
            worker_list (tuple(train_worker_list,test_worker_list)): 包含训练客户端列表和测试客户端列表
            model (torch.model): 实现了torch.model类的子类
            server_params (dict): 服务器-联邦学习相关的参数
            client_params (dict): 客户端-一次训练相关的参数
        """        
        # 内部参数---------------------------------------------------------------------------------------
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        
        
        # 模型优化器参数----------------------------------------------------------------------------------
        self.model = copy.deepcopy(model).to(self.device)
        # 优化器
        # self.is_aggregate_optimizer = False
        # 需要进行参数聚合的Adam优化器
        # self.optimizer = optim.Adam(self.model.parameters(),lr=0.001)
        # 不需要进行参数聚合的SGD优化器
        self.optimizer = torch.optim.SGD(self.model.parameters(),lr=0.01)
        # 带有正则化项的Adam优化器
        # self.optimizer = optim.Adam(self.model.parameters(),lr=0.001,weight_decay=0.000001)
        # 带有正则化项的SGD优化器
        # self.optimizer = torch.optim.SGD(self.model.parameters(),lr=0.01,weight_decay=0.000001)
        

        # 服务器策略控制参数---------------------------------------------------------------------------------
        # 通信的轮次
        if("rounds" in server_params):
            self.rounds=server_params["rounds"]
        else:
            self.rounds = 200
        # 客户端参与度。表示每轮有多少个客户端参与到训练当中
        if("ratio" in server_params):
            self.ratio = server_params["ratio"]
        else:
            self.ratio = 1.0
        
        # 本地更新的参数，这个参数就用epochs来表示吧。暂时保留
        if("inner_lr" in server_params):
            self.inner_lr = server_params["inner_lr"]
        else:
            self.inner_lr =0.01 
        if("outer_lr" in server_params):
            self.outer_lr = server_params["outer_lr"]
        else:
            self.outer_lr = 0.8

        # 根据innerlr修改optimizer
        optimizer_state_dict = self.optimizer.state_dict()
        optimizer_state_dict['param_groups'][0]['lr']=self.outer_lr
        self.optimizer.load_state_dict(optimizer_state_dict)


        # 客户端配置-----------------------------------------------------------------------------------------
        # 创建多个客户端列表.也可以根据数据集的数量自己确定训练节点和测试节点的数量。
        self.train_worker_list,self.test_worker_list = worker_list 
        # 训练节点的个数和测试节点的个数，应该由数据集的划分决定。数据集中给了多少组训练数据就有多少个训练节点。测试节点类似。
        # self.client_size=len(self.train_worker_list)这两个值可能被更改，例如随机选择指定数量的客户端
        # self.test_size=len(self.test_worker_list)这个值可能被更改，应该动态确定。
        # 加载参数。启动worker处理数据。之后可以改成异步。
        # 两种数据处理方式，都进行实现。不能鲁莽更改。应该还是具体情况按需实现。保持原样，暂时不动，下一轮重构再说。
        for worker in self.train_worker_list:
            worker.set_parameters(client_params)
            worker.process_data()
            # worker.process_data_support_and_query()
        for worker in self.test_worker_list:
            worker.set_parameters(client_params)
            worker.process_data()
            # worker.process_data_support_and_query()

        # 状态参数--------------------------------------------------------------------
        # 服务器迭代过程中，每一个round，参数都会更新，是服务器当前状态的描述。model_parameters本质上也是参数的状态。
        # 重构中考虑设计模式，避免在实现业务逻辑的是偶进行过度设计
        # 上一轮保留的模型的参数。这里需要保留一份copy，如果是单纯的赋值的话是引用，会出现问题
        self.old_model_parameters = copy.deepcopy(list(self.model.parameters()))
        # 每一轮，累计的客户端贡献度
        self.information_increase_matrix=torch.zeros([len(self.train_worker_list)])
        # 根据累计的客户端贡献度计算下一轮的合并权重
        self.client_weight = []
        # 每一轮，模型变化的剧烈程度
        self.information_increase_list = []

        # 每一轮随机挑选的客户端。需要不断的更新
        self.selected_train_worker_list=[]
        # 训练损失、训练准确率、测试准确率
        self.test_accuracy = []
        self.train_loss = []
        self.train_accuracy = []

        # 差分隐私参数--------------------------------------------------------------------------------------
        # 通过rayleigh.ppf函数。添加通信的噪声。即在发送信息之前，添加一个随机值。借鉴自scaffold的代码实现
        # self.communication_thresh = 0
        # if self.noise:
        #     self.communication_thresh = rayleigh.ppf(1 - users_per_round / total_users)  # h_min

        logger.info("FedServerBase init: device={}  rounds={} client_size={} test_size={}".format(self.device,self.rounds,len(self.train_worker_list),len(self.test_worker_list)))


    # 计算两组值之间的方差。表示变化率，表示两个向量的差的范数
    def get_params_norm(self,old_param_list,new_param_list):
        """计算两组值之间的方差。表示变化率或者说贡献率。这个是工具函数。与类没哟耦合

        Args:
            old_param_list (list[tensor]): 原有的参数
            new_param_list (list[tensor]): 最新的参数

        Returns:
            float: 梯度的二范数
        """  
        params = []

        for old_param,new_param in zip(old_param_list,new_param_list):
            params.append(torch.flatten(new_param  - old_param))

        # return torch.linalg.norm(torch.cat(params), 2)
        return float(torch.norm(torch.cat(params)))

    def get_tensor_vector_from_dict(self,param_dict):
        """将高维的tensor组成的字典。转换成一维的tensor。方便计算。

        Args:
            param_dict (dict[tensor]): 高维tensor 字典

        Returns:
            tensor: 一维的tensor
        """        
        # 首先吧字典内部的所有向量展开并串联起来
        params = []
        for name in param_dict:
            params.append(torch.flatten(param_dict[name]))
        
        # 把tensor_list转换为tensor_vector
        return torch.cat(params)
    
    def get_tensor_vector_from_list(self,param_list):
        """将高维的tensor组成的列表，转换成一维tensor

        Args:
            param_list (list[tensor]): 高维tensor 列表

        Returns:
            tensor: 一维的tensor vector
        """        
        # 首先吧字典内部的所有向量展开并串联起来
        params = []
        for param in param_list:
            params.append(torch.flatten(param))
        
        # 把tensor_list转换为tensor_vector
        return torch.cat(params)
    
    # TODO:是否将这个作为单独的模块呢？是否将差分隐私也作为一个单独模块呢？
    # 我觉得暂时没有必要将它涉及为单独的模块。任务量的增加，对工程进行重构。
    # 然后将这一部分集成的功能独立出去，作为功能扩展。正好，花时间使用C++来实现一下涉及模式之后
    # 进行重构，将贡献度计算模块和差分隐私模块独立出去。
    def calculate_imformation_increase(self,model_state_dict_list,old_accuracy,new_accuracy):
        """计算训练过程中的信息增益速度。用来衡量贡献度。待会进行计算。考虑到计算的准确率\n
        1. 计算梯度方向的投影
        2. 将所有客户端的投影进行归一化。得到本次贡献度的比例。a = x 內积 y /|y|
        3. 将accuracy变化作为本次round的权重。乘以归一化后的值进行保存。
        4. 为每一个客户端维护一个数值，表示当前累计的贡献度。形成一个向量。
        5. 返回本轮的贡献度比例向量。
        Args:
            old_param (list[tensor]): 原来参数
            new_param (list[tensor]): 本次参数
            accuracy (float): 准确率
        """
        with torch.no_grad():    
            # 因为存储的参数是一个列表形式，需要展开转换为1维的tensor，然后晕眩
            new_model_parameters = self.model.parameters()
            old_model_parameters_vector = self.get_tensor_vector_from_list(self.old_model_parameters)
            new_model_parameters_vector = self.get_tensor_vector_from_list(new_model_parameters)
            delta_model_parameters_vector = new_model_parameters_vector-old_model_parameters_vector
            # 求单位变化方向的向量.除以模长
            delta_param_vector  = delta_model_parameters_vector/ torch.norm(delta_model_parameters_vector,2)
    
            self.information_increase_list.append( torch.norm(delta_model_parameters_vector,2).item())

            information = torch.zeros([len(model_state_dict_list)])
            # 第i个客户端
            for i,model_state_dict in enumerate(model_state_dict_list):
                # 计算模型参数的变更方向，并将所有的内容展开，连接为一个向量
                client_model_parameters_vector = self.get_tensor_vector_from_dict(model_state_dict)
                client_delta_parameters_vector = client_model_parameters_vector-old_model_parameters_vector
                # 得到第i个客户端的信息增益
                information[i] = (delta_param_vector* client_delta_parameters_vector).sum()

            # 将信息增益矩阵进行百分比化.并乘以当前的变量。
            information =  information/information.sum()*(new_accuracy-old_accuracy)

            # 将信息增益矩阵增加准确率变化的考量
            self.information_increase_matrix += information
        # 返回当前轮次的信息增益矩阵
        return information
    # 返回当前的累计贡献度向量
    def get_information_increase(self):
        return self.information_increase_matrix

    # 根据累计的贡献量矩阵，计算权重
    def calculate_client_weight(self):
        """
        1. 首先将信息矩阵归一化
        2. 然后添加指数影响因子，全部为正，且缩小影响差距。

        Returns:
            client_weight: 权重列表
        """        
        # if(torch.sum(self.information_increase_matrix)>1):
        #     temp_information_metrix = self.information_increase_matrix/torch.sum(self.information_increase_matrix)
        temp_weight = 1/(1+torch.exp(-self.information_increase_matrix))
        self.client_weight = temp_weight/temp_weight.sum()
        return self.client_weight
    #TODO: 应用差分依隐私的内容。以后将这一部分独立出去。
    def apply_diffrential_privacy(self):
        pass
    
    # 应用客户端选择方案。当前是随机选择
    def random_select_clients(self):
        import random
        # 计算应该选择的客户端
        number = self.ratio*len(self.train_worker_list)
        self.selected_train_worker_list = random.sample(self.train_worker_list,number)

    def aggregate_by_weight(self,model_state_dict_list):   
        """基于权重的参数聚合方案。这里以fedavg算法的权重聚合为模板。

        Args:
            model_state_dict_list (dict[tensor]): pytorch格式的模型权重
        """             
        # fed_avg算法
        new_par = copy.deepcopy(self.model.state_dict())
        pi = torch.tensor(1.0/len(model_state_dict_list))
        for name in new_par:
            new_par[name] = torch.zeros(new_par[name].shape).to(self.device)
        # 使用联邦平均算法进行聚合
        for model_dict in model_state_dict_list:
            for name in new_par:
                new_par[name]+= model_dict[name]*pi
        # 将聚合后的参数加载到模型当中
        self.model.load_state_dict(copy.deepcopy(new_par))

    
    # 尝试自己写梯度下降算法
    def aggregate_by_grad(self,grad_dict_list):
        """基于梯度的权重聚合方案

        Args:
            grad_dict_list (list[tensor]): pytorch格式的梯度
        """        
      
        # 定义0梯度字典
        new_grad_dict = copy.deepcopy(self.model.state_dict())
        pi = torch.tensor(1.0/len(grad_dict_list))
        for name in new_grad_dict:
            new_grad_dict[name] = torch.zeros(new_grad_dict[name].shape).to(self.device)
        
        # 将梯度字典累加，构成最后的梯度
        for grad_dict in grad_dict_list:
            for name in new_grad_dict:
                new_grad_dict[name]+= grad_dict[name]*pi

        # 将聚合后的梯度放到model_state_dict当中。累加之后进行梯度下降。可以在这里尝试使用优化器。optimizer
        model_state_dict = self.model.state_dict()
        for name in new_grad_dict:
            model_state_dict[name]-=new_grad_dict[name]
        # 将聚合后的参数加载到模型当中
        self.model.load_state_dict(model_state_dict)


    
    def aggregate_optimizer(self,optim_state_dict_list):
        """优化器聚合方案。主要对应于原来使用了adam优化器，进行参数聚合的方案。\n
        后来这个方法弃用了。

        Args:
            optim_state_dict_list (dict[tensor]): pytorch格式的优化器参数
        """        
        new_opt_par = copy.deepcopy(optim_state_dict_list[0])
        #print(new_opt_par["state"]) #这是每层神经网络的参数:step,vt,mt,是一个Dict
        #print(new_opt_par["param_groups"][0][ 'params'])#这是模型每一层参数的索引，是一个list
        #print(new_opt_par["param_groups"][0])#这是模型每一层参数的索引，是一个list,state param_groups
        
        # 将所有的参数初始化为0
        for index in new_opt_par["param_groups"][0][ 'params']:

            new_opt_par["state"][index]['step'] = 0
            new_opt_par["state"][index]['exp_avg'] = torch.zeros(new_opt_par["state"][index]['exp_avg'].shape)
            new_opt_par["state"][index]['exp_avg_sq'] = torch.zeros(new_opt_par["state"][index]['exp_avg_sq'].shape)
            #print(new_opt_par["state"][index]) #这是index这一层神经网络的参数:step,vt,mt,是一个Dict

        for idx, par in enumerate(optim_state_dict_list):
            pi = torch.tensor(1.0/len(optim_state_dict_list))
            for i in range(len(new_opt_par["param_groups"][0][ 'params'])):
                index = new_opt_par["param_groups"][0]['params'][i] #第i层参数的索引
                index2 = par["param_groups"][0]['params'][i] #第i层参数的索引
                new_opt_par["state"][index]['step'] += par["state"][index2]['step'] * pi 
                new_opt_par["state"][index]['exp_avg'] += par["state"][index2]['exp_avg']* pi 
                new_opt_par["state"][index]['exp_avg_sq'] += par["state"][index2]['exp_avg_sq']* pi 

        #print(new_opt_par["state"]) #这是每层神经网络的参数:step,vt,mt,是一个Dict
        self.optimizer.load_state_dict(copy.deepcopy(new_opt_par))

    # 定义了训练的整个过程。
    def start_federated(self):
        """训练过程。可以重复调用，延续之前的训练。但是不支持修改参数继续训练。
        1. 客户端加载当前模型
        2. 客户端执行训练返回loss和accuracy
        3. 客户端返回模型参数
        4. 聚合客户端模型参数
        5. 测试模型，并输出或者保存结果。
        """        
        # 客户端进行训练 迭代以下过程
        for round in range(self.rounds):
            results_model = []
            # results_optim = []
            running_accuracy = []

            # 发送模型。为每一个客户端发送一组新的模型复制
            for worker in self.train_worker_list:
                worker.load_model_state_dict(copy.deepcopy(self.model.state_dict()))
            
            # 训练模型。经协程的折磨。我发现必须底层支持协程，否则这个东西本质上是不支持协程的。
            # 所以使用多线程。建立与worker的关系。在多线程中执行训练。
            result_thread = []
            for worker in self.train_worker_list:
                # running_accuracy.append(worker.train())
                # 尝试多线程版本
                t=WorkerThread(worker.train,)
                t.start()
                result_thread.append(t)
     
            for t in result_thread:
                # print(t.get_result())
                running_accuracy.append(t.get_result())
    
            # 返回训练参数
            for worker in self.train_worker_list:
                results_model.append(worker.get_model_state_dict())
                # results_optim.append(worker.get_optim_state_dict())
            
            # 聚合训练参数。聚合策略的选择
            self.aggregate_by_weight(results_model)
            
            # 对模型进行验证
            result_test=0
            for worker in self.test_worker_list:
                worker.load_model_state_dict(copy.deepcopy(self.model.state_dict()))
                result_test += worker.test()[0]/len(self.test_worker_list)

            # 数据记录(计算信息矩阵、保存测试准确率、训练损失、训练准确率、当前的模型参数保存为旧模型的参数)
            if round>0:
                self.calculate_imformation_increase(results_model,self.test_accuracy[-1],result_test)

            self.test_accuracy.append(result_test)
            running_state = torch.tensor(running_accuracy)
            running_state = torch.mean(running_state,dim=0)
            self.train_loss.append(running_state[0].item())
            self.train_accuracy.append(running_state[1].item())
            self.old_model_parameters = copy.deepcopy(list(self.model.parameters()))
            logger.info("FedServerBase start_federated:round={},test_acc={:.2f},train_loss={:.2f},train_acc={:.2f}".format(round, result_test,running_state[0].item(),running_state[1].item()))

    def save_model(self,model_path):
        """持久化模型

        Args:
            model_path (str): 模型保存路径
        """        
        torch.save(self.model, model_path)

    def read_model(self,model_path):
        """加载模型

        Args:
            model_path (str): 模型加载路径
        """        
        self.model = torch.load(model_path)

    # 通过torch将内部状态保存到指定的向量当中。
    # 为什么不用numpy。因为没哟导入numpy。而且既然torch能做到。而且本模块引入了torch。那么就可以用torch保存
    def save_state(self,path):
        """保存内部的状态变量。方便进行数据分析。通过字典的形式保存

        Args:
            path (str):存储文件的路径和名字 
        """        
        state_dict = {
            'information_increase_matrix':self.information_increase_matrix,
            'information_increase_list':self.information_increase_list,
            'test_accuracy':self.test_accuracy,
            'train_loss':self.train_loss,
            'train_accuracy':self.train_accuracy,
        }
        torch.save(state_dict,path)

    #TODO： 运行过程中打点，通过第三方工具，展示运行的画面
    def metrics(self):
        pass
class optimizer(torch.optim.Optimizer):
    """可以定义自己的优化器。实现梯度下降过程。最好定义一个新文件，来做这一件事情。

    Args:
        torch (torch.Optimizer): torch提供的优化器的基类。可以重写step方法。
    """    
    pass



# 实现了自己的多线程类

import threading


class WorkerThread(threading.Thread):

    def __init__(self,func,args=()):
        super(WorkerThread,self).__init__()
        self.func = func
        self.args = args

    def run(self):
        self.result = self.func(*self.args)

    def get_result(self):
        threading.Thread.join(self) # 等待线程执行完毕
        try:
            return self.result
        except Exception:
            return None
